/**
 * Copyright 2024-2024 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "model_builder_ndk_impl.h"

#include <climits>
#include <dlfcn.h>

#include "securec.h"
#include "infra/base/assertion.h"
#include "infra/base/securestl.h"

#include "framework/infra/log/log.h"
#include "base/common/file_util/file_util.h"
#include "model/built_model/customdata_util.h"

#include "model_builder/om/model_build_ndk_options_util.h"
#include "model/built_model/built_model_ndk_impl.h"
#include "model_manager/general_model_manager/core/model_manager_ndk_impl.h"
#include "model_manager/general_model_manager/ndk/ndk_util/ndk_util.h"
#include "model_manager/general_model_manager/ndk/hiai_ndk/hiai_ndk_nncore.h"

namespace hiai {
namespace {
Status BuildModel(const ModelBuildOptions& options, const std::string& modelName,
    const std::shared_ptr<IBuffer>& modelBuffer, std::shared_ptr<IBuiltModel>& builtModel)
{
    OH_NNCompilation* nnCompilation = HIAI_NDK_NNCompilation_ConstructWithOfflineModelBuffer(
        modelBuffer->GetData(), modelBuffer->GetSize());
    HIAI_EXPECT_NOT_NULL(nnCompilation);

    // 向compilation中设置buildoption
    Status ret = ModelBuildNDKOptionsUtil::AssignToNNCompilation(nnCompilation, options);
    if (ret != SUCCESS) {
        HIAI_NDK_NNCompilation_Destroy(&nnCompilation);
        FMK_LOGE("AssignToNNCompilation failed.");
        return FAILURE;
    }

    // 执行编译
    ret = HIAI_NDK_NNCompilation_Build(nnCompilation);
    if (ret != SUCCESS) {
        HIAI_NDK_NNCompilation_Destroy(&nnCompilation);
        FMK_LOGE("HIAI_NDK_NNCompilation_Build failed.");
        return FAILURE;
    }
    builtModel = make_shared_nothrow<BuiltModelNDKImpl>(std::shared_ptr<OH_NNCompilation>(
        nnCompilation,
        [](OH_NNCompilation* p) {
            HIAI_NDK_NNCompilation_Destroy(&p);
        }),
        std::const_pointer_cast<IBuffer>(modelBuffer));
    if (builtModel == nullptr) {
        FMK_LOGI("build make_shared failed.");
        HIAI_NDK_NNCompilation_Destroy(&nnCompilation);
        return FAILURE;
    }
    builtModel->SetName(modelName);
    FMK_LOGI("build model success.");
    return SUCCESS;
}
}

Status ModelBuilderNDKImpl::Build(const ModelBuildOptions& options, const std::string& modelName,
    const std::shared_ptr<IBuffer>& modelBuffer, std::shared_ptr<IBuiltModel>& builtModel)
{
    if (modelBuffer == nullptr) {
        FMK_LOGE("modelBuffer is nullptr.");
        return INVALID_PARAM;
    }

    if (modelName.length() > PATH_MAX) {
        FMK_LOGE("modelName length is too long.");
        return INVALID_PARAM;
    }

#ifdef AI_SUPPORT_AIPP_API
    std::shared_ptr<IBuffer> compatibleModelBuffer = nullptr;
    CustomModelData customData;
    if (CustomDataUtil::MakeCompatibleBuffer(compatibleModelBuffer, customData, modelBuffer) != SUCCESS) {
        return FAILURE;
    }

    if (compatibleModelBuffer != nullptr) {
        if (BuildModel(options, modelName, compatibleModelBuffer, builtModel) != SUCCESS) {
            return FAILURE;
        }
        if (builtModel != nullptr) {
            builtModel->SetCustomData(customData);
        }
        return SUCCESS;
    }
#endif
    return BuildModel(options, modelName, modelBuffer, builtModel);
}

Status ModelBuilderNDKImpl::Build(const ModelBuildOptions& options, const std::string& modelName,
    const std::string& modelFile, std::shared_ptr<IBuiltModel>& builtModel)
{
    buffer_ = FileUtil::LoadToBuffer(modelFile);
    if (buffer_ == nullptr) {
        return FAILURE;
    }

    std::shared_ptr<IBuffer> localBuffer =
        CreateLocalBuffer(static_cast<void*>(buffer_->MutableData()), buffer_->GetSize(), false);

    return Build(options, modelName, localBuffer, builtModel);
}

std::shared_ptr<IModelBuilder> CreateModelBuilderFromNDK()
{
    return make_shared_nothrow<ModelBuilderNDKImpl>();
}

} // namespace hiai
